import torch
from torch import nn, Tensor
from typing import Tuple
from torch.nn import functional as F
import warnings
import math

from .backbone import resnet18, resnet34, resnet50, resnet101, resnet152,mobilenet,xception
from .head import Deeplabv3plusHead
from .builder import EncoderDecoder


class DeepLabv3plus(EncoderDecoder):
    def __init__(self, variant: str = 'resnet50', num_classes: int = 19, pretrained: bool = False):

        backbone = eval(variant)(pretrained=pretrained)
        if variant == 'resnet18' or variant == 'resnet34':
            in_channels = 512
            low_level_channels = 128
        elif variant == 'resnet50' or variant == 'resnet101':
            in_channels = 2048
            low_level_channels = 512
        elif variant=="xception":
            in_channels = 2048
            low_level_channels = 256
        elif variant=="mobilenet":
            in_channels = 320
            low_level_channels = 24

        else:
            raise ValueError('Unsupported backbone - `{}`, Use mobilenet, xception.'.format(backbone))


        decode_head = Deeplabv3plusHead(in_channels=in_channels, low_level_channels=low_level_channels,
                                        num_classes=num_classes)
        super(DeepLabv3plus, self).__init__(backbone=backbone, decode_head=decode_head,
                                            in_channels=[low_level_channels, in_channels])


if __name__ == '__main__':
    x = torch.randn(2, 3, 224, 224)
    model = DeepLabv3plus(variant="mobilenet", num_classes=19)
    y = model(x)
    print(y.shape)
